#!/usr/bin/env python3

"""
@module  unbricker
@version 0.0.1
@date    2024.06.18
@author  Alexander Zhevak
@email   zhevak@mail.ru

@brief   Раскирпичиватель CH32V003J4M6
"""

import usb.core
import usb.util
import sys


VID          = 0x1A86  # VID
PID_RV       = 0x8010  # PID в режиме RISC-V
PID_ARM      = 0x8012  # PID в режиме ARM
EP_CMD       = 0x01    # Адрес конечной точки для передачи команд в программатор в режиме RISC-V
EP_ANS       = 0x81    # Адрес конечной точки для приёма ответов из программатора в режиме RISC-V
EP_CMD_ARM   = 0x01    # Адрес конечной точки для передачи команд в программатор в режиме ARM
PACKET_SIZE  = 64      # Размер пакетов


class Exc(Exception):
    def __init__(self, msg):
        self.msg = msg
       

class WCHLinkE:

    device = None

    def __init__(self):
        ''' Подключаемся к программатору '''
    
        # Предполагаю, что программатор находится в режиме RISC-V
        self.device = usb.core.find(idVendor = VID, idProduct = PID_RV)
        if self.device is None:
            # Программатор, работающий в режиме RISC-V, не обнаружен
            # Предполагаю, что программатор находится в режиме ARM
            self.device = usb.core.find(idVendor = VID, idProduct = PID_ARM)
            if self.device:
                # Программатор работает в режиме ARM
                self.device = self.goto_riscv()  # Перевожу в режим RISC-V
                self.device = usb.core.find(idVendor = VID, idProduct = PID_RV)

        if self.device is None:
            raise Exc('Программатор WCH-Link не обнаружен')


    def goto_riscv(self):
        ''' Переводит программатор в режим работы RISC-V '''
  
        self.device.write(EP_CMD_ARM, b'\x81\xff\x01\x52')
    
        counter = 30
        while counter:
            time.sleep(0.1)
            counter -= 1
            self.device = usb.core.find(idVendor = VID, idProduct = PID_RV)
            if self.device:
                return
                
        raise Exc('Не удалось переключить программатор в режим RISC-V')


    def identify(self):
        ''' Проверяет, что программатор WCHLink-E, а не какой-то другой '''

        # Получаю информацию о программаторе
        self.device.write(EP_CMD, b'\x81\x0d\x01\x01')
        reply = self.device.read(EP_ANS, PACKET_SIZE, 3000)
    
        if reply[5] != 0x12:
            raise Exc('Это какой-то другой программатор, а не WCHLink-E')


    def clear_buffer(self):
        ''' Очищает буфер ответов '''

        try:
            self.device.read(EP_ANS, PACKET_SIZE, 1)
        except:
            pass


    def unbrick(self):
        ''' Производит раскирпичивание '''

        self.device.write(EP_CMD, b'\x81\x0d\x01\x0f\x09')
        answer = self.device.read(EP_ANS, PACKET_SIZE, 3000)

        if answer[3] != 15:
            raise Exc('После процесса раскирпичивания микроконтроллер так и не появился.')
        

if __name__ == "__main__":

    try:
        wchlinke = WCHLinkE()        
        wchlinke.clear_buffer() # На всякий случай очищу буфер ответов
        wchlinke.identify()     # Убеждаюсь, что это программатор WCHLink-E
        wchlinke.unbrick()      # Раскирпичиваю
    except Exc as exc:
        print(exc)
        exit(1)

    print('Микроконтроллер сейчас должен быть вроде бы раскирпичен. Проверяйте!')

